#! /usr/bin/python3

import os
import xml.etree.cElementTree as ET

import numpy as np
from IPython import embed
from scipy.spatial.transform import Rotation

from .modern_robotics import ScrewToAxis
from .Utils import ToTransformNormForm
from .ZenLogger import ZenLogger

logger = ZenLogger.get_logger(__name__)

CWD = os.path.dirname(__file__)


class RobotTree:
    def __init__(self, path):
        self.m_urdf_path = path
        self.m_urdf_dir = os.path.dirname(self.m_urdf_path)
        self.m_tree = ET.parse(self.m_urdf_path)
        self.m_robot = self.m_tree.getroot()
        self.m_name = self.m_robot.get("name", "robot")
        self.joint_names = []
        self.link_names = []
        self.m_links = dict()
        self.m_joints = dict()
        self.m_dof = 0
        self.m_s_list = None
        self.__Init()

    def GetName(self):
        return self.m_name

    def GetDoF(self):
        return self.m_dof

    def __ComputeSList(self):
        q = np.zeros(3)  # a point at screw axis
        T = np.eye(4)
        axes = []
        for name in self.joint_names:
            joint = self.m_joints[name]
            if joint["type"] == "fixed":
                continue
            axis = joint["axis"]
            origin_xyz = joint["origin_xyz"]
            origin_rpy = joint["origin_rpy"]
            q = T @ np.r_[origin_xyz, 1]
            axes.append(ScrewToAxis(q[:3], axis, 0))
            T = T @ ToTransformNormForm(origin_xyz, origin_rpy)
        self.m_s_list = np.array(axes).T
        self.m_home_pose = T

    def __GetLinks(self):
        # TODO: check the sequence
        self.m_links = {}
        links = self.m_robot.findall("link")
        for l in links:
            link = {}
            link["name"] = l.get("name")
            link["parent"] = ""
            link["child"] = ""
            link["visual_mesh"] = ""
            link["collision_mesh"] = ""
            if link["name"] != "ee_link":
                link_visual = l.find("visual")
                if link_visual:
                    visual_file = link_visual.find(
                        "geometry").find("mesh").get("filename").split("//")[1]
                    link["visual_mesh"] = os.path.join(
                        self.m_urdf_dir, visual_file)
                    rgbastring = link_visual.find(
                        "material").find("color").get("rgba")
                    link["color"] = np.fromstring(
                        rgbastring, dtype=float, sep=" ")
                    link_collision = l.find("collision")
                if link_collision:
                    collision_file = link_collision.find(
                        "geometry").find("mesh").get("filename").split("//")[1]
                    link["collision_mesh"] = os.path.join(
                        self.m_urdf_dir, collision_file)
            self.m_links[link["name"]] = link

    def __GetJoints(self):
        # TODO: check the sequence
        self.m_joints = {}
        self.m_dof = 0
        joints = self.m_robot.findall("joint")
        active_joints = []
        active_links = []
        active_joints_axis = []
        for j in joints:
            joint = {"limit": {}}
            joint["name"] = j.get("name")
            joint["type"] = j.get("type")
            joint["parent"] = j.find("parent").get("link")
            joint["child"] = j.find("child").get("link")
            self.m_links[joint["child"]]["parent"] = joint["name"]
            self.m_links[joint["parent"]]["child"] = joint["name"]
            if joint["type"] != "fixed":
                self.m_dof += 1
                joint["origin_xyz"] = np.fromstring(
                    j.find("origin").get("xyz"), dtype=float, sep=" ")
                joint["origin_rpy"] = np.fromstring(
                    j.find("origin").get("rpy"), dtype=float, sep=" ")
                joint["axis"] = np.fromstring(
                    j.find("axis").get("xyz"), dtype=float, sep=" ")
                joint["limit"]["effort"] = float(
                    j.find("limit").get("effort", 0))
                joint["limit"]["upper"] = float(
                    j.find("limit").get("upper", 3))
                joint["limit"]["lower"] = float(
                    j.find("limit").get("lower", -3))
                joint["limit"]["velocity"] = float(
                    j.find("limit").get("velocity", 1))
                joint["limit"]["acceleration"] = float(
                    j.find("limit").get("acceleration", 5))
                active_joints.append(joint['name'])
                active_joints_axis.append(joint['axis'])
            self.m_joints[joint["name"]] = joint
        self.active_joints = tuple(active_joints)
        self.active_links = tuple([self.m_joints[j]['child']
                                   for j in self.active_joints])
        self.active_joints_axis = tuple(active_joints_axis)

    def __Init(self):
        self.__GetLinks()
        self.__GetJoints()
        for l in self.m_links.values():
            if not l["parent"]:
                self.m_root_link = l["name"]
        self.GetChain()
        self.__ComputeSList()
        logger.info("Robot chain: {}".format(self.GetChain()))
        logger.info("Robot links: {}".format(self.link_names))
        logger.info("Robot joints: {}".format(self.joint_names))

    def GetRootLink(self):
        return self.m_root_link

    def GetLinks(self):
        return self.m_links

    def GetJoints(self):
        return self.m_joints

    def GetChain(self):
        chain = []
        node = self.m_links[self.GetRootLink()]
        self.joint_names = []
        self.link_names = []
        while node["child"]:
            chain.append(node["name"])
            child = node["child"]
            node = self.m_joints.get(child, False)
            if node:
                self.joint_names.append(node["name"])
            else:
                node = self.m_links.get(child, False)
                if node:
                    self.link_names.append(node["name"])
                else:
                    logger.error("node child {} not found!".format(child))
                    break
        chain.append(node["name"])
        return chain


if __name__ == "__main__":
    CWD = os.path.dirname(__file__)
    upper_dir = os.path.dirname(CWD)
    # filename = os.path.join(upper_dir, "Models/GP12/GP12.urdf")
    filename = os.path.join(upper_dir, "models/Motomini/Motomini.urdf")
    tree = RobotTree(filename)
    embed()
    # print(tree.GetLinks())
    # print(tree.GetJoints())
    print("{} root: {}".format(tree.GetName(), tree.GetRootLink()))
    # print(tree.GetChain())
